import pandas as pd
from typing import Dict, List, Literal
from io import StringIO


def read_excel(excel_path: str, dtype=str) -> Dict[str, pd.DataFrame]:
    """Read excel file"""
    excel_data = pd.read_excel(excel_path, sheet_name=None, dtype=dtype)
    if dtype is str:
        for df in excel_data.values():
            df.fillna("", inplace=True)
    return excel_data


def read_csv(csv_path: str, dtype=str) -> pd.DataFrame:
    """Decode csv file by utf-8 or gb18030 and read it"""
    for encoding in ['utf-8', 'gb18030']:
        try:
            df = pd.read_csv(csv_path, dtype=dtype, encoding=encoding)
            if dtype is str:
                df.fillna("", inplace=True)
            return df
        except UnicodeDecodeError:
            pass
    raise RuntimeError('File ' + csv_path + ' cannot be decoded by utf-8 or gb18030')


def read_json(json_path: str, dtype=str) -> pd.DataFrame:
    """Decode json file by utf-8 or gb18030 and read it"""
    for encoding in ['utf-8', 'gb18030']:
        try:
            df = pd.read_json(json_path, dtype=dtype, encoding=encoding)
            if dtype is str:
                df.fillna("", inplace=True)
            return df
        except UnicodeDecodeError:
            pass
    raise RuntimeError('File ' + json_path + ' cannot be decoded by utf-8 or gb18030')


def strdata2df(csv_data: str, dtype=str) -> pd.DataFrame:
    """csv string to df"""
    df = pd.read_csv(StringIO(csv_data), dtype=dtype)
    if dtype is str:
        df.fillna("", inplace=True)
    return df


def equal(df1: pd.DataFrame, df2: pd.DataFrame) -> bool:
    return df1.to_dict() == df2.to_dict()


def reset_index(df: pd.DataFrame) -> pd.DataFrame:
    return df.reset_index(drop=True)


def dropna(df: pd.DataFrame) -> pd.DataFrame:
    return df.dropna().reset_index(drop=True)


def drop_duplicates(
        df: pd.DataFrame,
        subset: List[str]=None,
        keep: Literal['first', 'last', False]='first') -> pd.DataFrame:
    if subset is None:
        subset = list(df.columns)
    return df.drop_duplicates(subset=subset, keep=keep).reset_index(drop=True)


def sort_values(df: pd.DataFrame, subset: List[str]=None) -> pd.DataFrame:
    if subset is None:
        subset = list(df.columns)
    return df.sort_values(by=subset).reset_index(drop=True)


def insert_row_at_first(df: pd.DataFrame, row: dict):
    return pd.concat([pd.DataFrame([row]), df], ignore_index=True)


def insert_column(df: pd.DataFrame, key: str, column: pd.Series, index=None) -> pd.DataFrame:
    if index is None:
        index = len(df.columns)
    df = df.reindex(columns=df.columns.insert(index, key))
    df[key] = column
    return df


def insert_column_before(df: pd.DataFrame, key: str, column: pd.Series, before_key: str) -> pd.DataFrame:
    assert before_key in df.columns
    index = df.columns.get_loc(before_key)
    df = df.reindex(columns=df.columns.insert(index, key))
    df[key] = column
    return df


def insert_column_after(df: pd.DataFrame, key: str, column: pd.Series, after_key: str) -> pd.DataFrame:
    assert after_key in df.columns
    index = df.columns.get_loc(after_key) + 1
    df = df.reindex(columns=df.columns.insert(index, key))
    df[key] = column
    return df
